# heartbeat_monitor.py
import asyncio
from datetime import datetime, timedelta
from utils.log_util import logger
from tcp_server.tcp_clients import ClientManager


class HeartbeatMonitor:
    """心跳监控器 - 定期检测客户端心跳超时"""
    
    def __init__(self, db_session, timeout_seconds=90, check_interval=30):
        """
        Args:
            db_session: 数据库会话工厂
            timeout_seconds: 心跳超时时间（秒），默认90秒
            check_interval: 检查间隔（秒），默认30秒
        """
        self.db_session = db_session
        self.timeout_seconds = timeout_seconds
        self.check_interval = check_interval
        self.running = False
        self._task = None
    
    async def start(self):
        """启动心跳监控"""
        if self.running:
            logger.warning("心跳监控已在运行")
            return
        
        self.running = True
        self._task = asyncio.create_task(self._monitor_loop())
        logger.info(f"心跳监控已启动（超时: {self.timeout_seconds}秒，检查间隔: {self.check_interval}秒）")
    
    async def stop(self):
        """停止心跳监控"""
        self.running = False
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
        logger.info("心跳监控已停止")
    
    async def _monitor_loop(self):
        """监控循环"""
        while self.running:
            try:
                await self._check_timeouts()
            except Exception as e:
                logger.exception(f"心跳监控异常: {e}")
            
            await asyncio.sleep(self.check_interval)
    
    async def _check_timeouts(self):
        """检查超时的客户端"""
        now = datetime.now()
        timeout_threshold = now - timedelta(seconds=self.timeout_seconds)
        
        # 获取所有在线客户端
        clients = ClientManager.get_all_clients()
        
        for client in clients:
            # 检查心跳是否超时
            if client.last_heartbeat < timeout_threshold:
                logger.warning(
                    f"客户端心跳超时: {client.client_id} "
                    f"(最后心跳: {client.last_heartbeat.strftime('%Y-%m-%d %H:%M:%S')})"
                )
                
                # 标记为离线并断开连接
                await self._handle_timeout_client(client.client_id)
    
    async def _handle_timeout_client(self, client_id: str):
        """处理超时的客户端"""
        try:
            client = ClientManager.get_client_by_id(client_id)
            if not client:
                return
            
            logger.warning(
                f"客户端心跳超时: {client_id} "
                f"(最后心跳: {client.last_heartbeat.strftime('%Y-%m-%d %H:%M:%S')})"
            )
            
            # ✅ 心跳超时应该关闭连接，原因：
            # 1. Agent每5秒发送心跳，90秒没收到说明至少18次心跳都丢失了，连接可能真的有问题
            # 2. Agent有自动重连机制，关闭连接后会自动重连
            # 3. 保持无效连接会占用服务器资源（文件描述符、内存）
            # 4. TCP连接本身如果已经断开，保持引用也没有意义
            
            # 先关闭TCP连接（如果连接还在）
            try:
                if client.writer and not client.writer.is_closing():
                    logger.debug(f"关闭超时客户端的TCP连接: {client_id}")
                    client.writer.close()
                    try:
                        await asyncio.wait_for(client.writer.wait_closed(), timeout=1.0)
                    except asyncio.TimeoutError:
                        pass  # 超时也继续，不阻塞
                    except Exception:
                        pass  # 忽略关闭时的异常
            except Exception as e:
                logger.debug(f"关闭TCP连接时出错（可能已关闭）: {e}")
            
            # 从内存中移除并更新数据库
            await ClientManager.remove_client(client_id, self.db_session)
            
            logger.info(f"已移除超时客户端: {client_id}")
            
        except Exception as e:
            logger.error(f"处理超时客户端失败 {client_id}: {e}")


# 全局心跳监控器实例
heartbeat_monitor = None

def get_heartbeat_monitor(db_session):
    """获取心跳监控器实例"""
    global heartbeat_monitor
    if heartbeat_monitor is None:
        heartbeat_monitor = HeartbeatMonitor(db_session)
    return heartbeat_monitor

